package com.xk._02算法篇._02unionFind.union;

import com.xk._02算法篇._02unionFind.Student;
import com.xk.tools.Asserts;

import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * @description:
 * @author: xu
 * @date: 2022/10/5 12:36
 */
public class GenericUnionFind<V> {
    private Map<V, Node<V>> nodes = new HashMap<>();

    private static class Node<V> {
        V value;
        Node<V> parent = this;
        int rank = 1;
        Node(V value) {
            this.value = value;
        }
    }

    public void makeSet(V v) {
        if (nodes.containsKey(v)) return;
        nodes.put(v, new Node<>(v));
    }

    public V find(V v) {
        Node<V> node = findNode(v);
        return node == null ? null : node.value;
    }

    public void union(V v1, V v2) {
        Node<V> p1 = findNode(v1);
        Node<V> p2 = findNode(v2);
        if (p1 == null || p2 == null) return;
        if (Objects.equals(p1.value, p2.value)) return;
        if (p1.rank < p2.rank) {
            p1.parent = p2;
        } else if (p1.rank > p2.rank) {
            p2.parent = p1;
        } else {
            p1.parent = p2;
            p2.rank += 1;
        }
    }

    public boolean isSame(V v1, V v2) {
        return Objects.equals(find(v1), find(v2));
    }

    /**
     * 找到v的根节点
     * @param v
     * @return
     */
    private Node<V> findNode(V v) {
        Node<V> node = nodes.get(v);
        if (node == null) return null;
        while (!Objects.equals(node.value, node.parent.value)) {
            node.parent = node.parent.parent;
            node = node.parent;
        }
        return node;
    }


    public static void main(String[] args) {
        GenericUnionFind<Student> uf = new GenericUnionFind<>();
        Student stu1 = new Student(15, "jack");
        Student stu2 = new Student(16, "rose");
        Student stu3 = new Student(17, "clock");
        Student stu4 = new Student(18, "jerry");

        uf.makeSet(stu1);
        uf.makeSet(stu2);
        uf.makeSet(stu3);
        uf.makeSet(stu4);

        uf.union(stu1, stu2);
        uf.union(stu3, stu4);
        uf.union(stu1, stu4);


        Asserts.test(uf.isSame(stu1, stu2));
        Asserts.test(uf.isSame(stu3, stu4));
        Asserts.test(uf.isSame(stu3, stu2));
        Asserts.test(uf.isSame(stu3, stu1));

    }

}
